import json
import os
from typing import Dict, List, Optional, Tuple

class RustKnowledgeVisualizer:
    """Rust知识可视化代理，用于生成Rust难理解知识的可视化图表"""
    
    def __init__(self):
        """初始化知识可视化代理"""
        self.api_key = os.getenv("OPENAI_API_KEY", "")
        self.base_url = "https://api.openai.com/v1/chat/completions"
        
        # 如果没有API密钥，使用模拟数据
        if not self.api_key:
            self.use_mock_data = True
        else:
            self.use_mock_data = False
    
    def generate_knowledge_visualization(self, topic: str, chart_type: str = "") -> Dict[str, str]:
        """
        生成Rust主题的知识可视化
        
        Args:
            topic: Rust主题
            chart_type: 图表类型（可选）
            
        Returns:
            包含可视化数据的字典
        """
        # 如果使用模拟数据
        if self.use_mock_data:
            return self._generate_mock_visualization(topic, chart_type)
        
        try:
            # 准备请求
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }
            
            # 构建提示词
            system_prompt = """
            你是一个专业的Rust编程语言专家，擅长将复杂概念可视化。
            请针对用户提供的Rust主题，生成一个Mermaid图表代码，用于可视化这个概念。
            图表应该清晰、简洁，易于理解。
            请只返回Mermaid代码，不要包含任何解释。
            可以使用以下图表类型：flowchart, classDiagram, stateDiagram, erDiagram, journey, gantt, pie, gitgraph。
            """
            
            data = {
                "model": "gpt-3.5-turbo",
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"请为Rust的'{topic}'概念生成一个{chart_type if chart_type else '合适的'}Mermaid可视化图表"}
                ],
                "temperature": 0.5,
                "max_tokens": 500
            }
            
            # 发送请求
            response = requests.post(self.base_url, headers=headers, json=data)
            response.raise_for_status()
            
            result = response.json()
            mermaid_code = result["choices"][0]["message"]["content"]
            
            return {
                "topic": topic,
                "mermaid_code": mermaid_code,
                "chart_type": self._detect_chart_type(mermaid_code),
                "status": "success"
            }
            
        except Exception as e:
            # 如果API调用失败，回退到模拟数据
            return self._generate_mock_visualization(topic)
    
    def _generate_mock_visualization(self, topic: str, chart_type: str = "") -> Dict[str, str]:
        """
        生成模拟可视化数据
        
        Args:
            topic: Rust主题
            chart_type: 图表类型
            
        Returns:
            包含可视化数据的字典
        """
        topic_lower = topic.lower()
        
        # 根据图表类型和主题生成对应的Mermaid代码
        if chart_type == "classDiagram" or (not chart_type and ("trait" in topic_lower or "特征" in topic_lower)):
            return self._generate_class_diagram(topic, topic_lower)
        elif chart_type == "stateDiagram" or chart_type == "stateDiagram-v2":
            return self._generate_state_diagram(topic, topic_lower)
        elif chart_type == "erDiagram":
            return self._generate_er_diagram(topic, topic_lower)
        elif chart_type == "journey":
            return self._generate_journey_diagram(topic, topic_lower)
        else:  # 默认使用flowchart
            return self._generate_flowchart(topic, topic_lower, chart_type)
    
    def _generate_flowchart(self, topic: str, topic_lower: str, chart_type: str = "") -> Dict[str, str]:
        """生成流程图"""
        if "所有权" in topic_lower or "ownership" in topic_lower:
            mermaid_code = """
flowchart LR
    A[变量创建] --> B[所有权转移]
    B --> C[旧变量失效]
    B --> D[新变量有效]
    D --> E[变量离开作用域]
    E --> F[内存释放]
    G[复制] --> H[两个变量都有效]
    I[克隆] --> J[深拷贝数据]
    J --> K[两个独立变量]
            """
        elif "借用" in topic_lower or "borrowing" in topic_lower:
            mermaid_code = """
flowchart TD
    A[原始数据] --> B[不可变引用&]
    A --> C[可变引用&mut]
    B --> D[多个不可变引用]
    C --> E[单一可变引用]
    D --> F[只读访问]
    E --> G[读写访问]
    F --> H[数据安全]
    G --> I[数据安全]
            """
        elif "生命周期" in topic_lower or "lifetime" in topic_lower:
            mermaid_code = """
flowchart LR
    A[变量创建] --> B[引用生成]
    B --> C[生命周期开始]
    C --> D[引用使用]
    D --> E[生命周期检查]
    E --> F{引用有效?}
    F -->|是| G[继续使用]
    F -->|否| H[编译错误]
    G --> I[变量销毁]
    I --> J[生命周期结束]
            """
        elif "trait" in topic_lower or "特征" in topic_lower:
            mermaid_code = """
flowchart TD
    A[定义Trait] --> B[声明方法]
    B --> C[为类型实现]
    C --> D[类型1实现]
    C --> E[类型2实现]
    D --> F[具体行为1]
            """
        elif "模式匹配" in topic_lower or "pattern matching" in topic_lower:
            mermaid_code = """
flowchart TD
    A[match表达式] --> B{匹配模式1}
    A --> C{匹配模式2}
    A --> D{匹配模式3}
            """
        elif "错误处理" in topic_lower or "error handling" in topic_lower:
            mermaid_code = """
flowchart TD
    A[可能失败的操作] --> B{操作成功?}
    B -->|是| C[返回Ok(T)]
            """
        elif "闭包" in topic_lower or "closure" in topic_lower:
            mermaid_code = """
flowchart LR
    A[闭包定义] --> B[捕获环境变量]
            """
        elif "并发" in topic_lower or "concurrency" in topic_lower:
            mermaid_code = """
flowchart TD
    A[并发编程] --> B[消息传递]
            """
        elif "宏" in topic_lower or "macro" in topic_lower:
            mermaid_code = """
flowchart TD
    A[宏调用] --> B[宏展开]
            """
        elif "迭代器" in topic_lower or "iterator" in topic_lower:
            mermaid_code = """
flowchart LR
    A[数据集合] --> B[创建迭代器]
            """
        else:
            # 默认通用图表
            mermaid_code = """
flowchart TD
    A[Rust概念] --> B[学习理解]
            """
        
        return {
            "topic": topic,
            "mermaid_code": mermaid_code,
            "chart_type": "flowchart",
            "status": "success"
        }
    
    def _generate_class_diagram(self, topic: str, topic_lower: str) -> Dict[str, str]:
        """生成类图"""
        if "所有权" in topic_lower or "ownership" in topic_lower:
            mermaid_code = """
classDiagram
    class Variable {
        +name: String
        +value: T
        +owner: Owner
        +create()
        +transfer()
        +drop()
    }
    class Owner {
        +id: String
        +variables: Vec~Variable~
        +transfer_ownership()
    }
    Variable --> Owner : owned by
    Owner o-- Variable : owns
            """
        elif "trait" in topic_lower or "特征" in topic_lower:
            mermaid_code = """
classDiagram
    class Trait {
        +方法签名
        +默认实现
    }
    class Type1 {
        +实现Trait方法
    }
    class Type2 {
        +实现Trait方法
    }
    Trait <|-- Type1
    Trait <|-- Type2
    note for Trait "定义共享行为"
    note for Type1 "具体实现1"
    note for Type2 "具体实现2"
            """
        elif "借用" in topic_lower or "borrowing" in topic_lower:
            mermaid_code = """
classDiagram
    class Data {
        +value: i32
    }
    class BorrowedRef {
        +data: &Data
        +borrow_immutable()
    }
    class MutRef {
        +data: &mut Data
        +borrow_mutable()
    }
    Data o-- BorrowedRef : immutable
    Data o-- MutRef : mutable
            """
        else:
            # 默认类图
            mermaid_code = """
classDiagram
    class RustConcept {
        +name: String
        +description: String
        +understand()
        +implement()
    }
    class Learner {
        +level: String
        +learn(concept)
    }
    RustConcept <-- Learner : learns
            """
        
        return {
            "topic": topic,
            "mermaid_code": mermaid_code,
            "chart_type": "classDiagram",
            "status": "success"
        }
    
    def _generate_state_diagram(self, topic: str, topic_lower: str) -> Dict[str, str]:
        """生成状态图"""
        if "所有权" in topic_lower or "ownership" in topic_lower:
            mermaid_code = """
stateDiagram-v2
    [*] --> VariableCreated
    VariableCreated --> OwnershipActive : 正常使用
    OwnershipActive --> OwnershipTransferred : 转移所有权
    OwnershipTransferred --> OldVariableDropped : 旧变量失效
    OwnershipTransferred --> NewOwnerActive : 新变量生效
    NewOwnerActive --> VariableDropped : 离开作用域
    OldVariableDropped --> MemoryDeallocated : 内存释放
    VariableDropped --> MemoryDeallocated
    MemoryDeallocated --> [*]
            """
        elif "借用" in topic_lower or "borrowing" in topic_lower:
            mermaid_code = """
stateDiagram-v2
    [*] --> DataCreated
    DataCreated --> ImmutableBorrow : 创建不可变引用
    DataCreated --> MutableBorrow : 创建可变引用
    ImmutableBorrow --> MultipleImmutable : 多个不可变引用
    MutableBorrow --> SingleMutable : 单一可变引用
    MultipleImmutable --> BorrowActive : 活跃借用
    SingleMutable --> BorrowActive : 活跃借用
    BorrowActive --> BorrowEnded : 引用结束
    BorrowEnded --> DataDropped
    DataDropped --> [*]
            """
        else:
            # 默认状态图
            mermaid_code = """
stateDiagram-v2
    [*] --> Start
    Start --> Learning : 开始学习
    Learning --> Understanding : 理解概念
    Understanding --> Practice : 实践应用
    Practice --> Mastery : 掌握概念
    Mastery --> [*]
            """
        
        return {
            "topic": topic,
            "mermaid_code": mermaid_code,
            "chart_type": "stateDiagram-v2",
            "status": "success"
        }
    
    def _generate_er_diagram(self, topic: str, topic_lower: str) -> Dict[str, str]:
        """生成实体关系图"""
        if "所有权" in topic_lower or "ownership" in topic_lower:
            mermaid_code = """
erDiagram
    VARIABLE {
        string name
        string type
        string value
    }
    OWNER {
        string owner_id
        string scope
        datetime created_at
    }
    TRANSFER {
        string transfer_id
        datetime transfer_time
        string from_owner
        string to_owner
    }
    VARIABLE ||--|| OWNER : owned_by
    OWNER ||--o{ TRANSFER : performs
    VARIABLE ||--o{ TRANSFER : transferred
            """
        elif "trait" in topic_lower or "特征" in topic_lower:
            mermaid_code = """
erDiagram
    TRAIT {
        string trait_name
        string method_signature
        string default_impl
    }
    TYPE {
        string type_name
        string type_params
    }
    IMPLEMENTATION {
        string impl_id
        string trait_name
        string type_name
        string method_body
    }
    TRAIT ||--o{ IMPLEMENTATION : defines
    TYPE ||--o{ IMPLEMENTATION : implements
            """
        else:
            # 默认ER图
            mermaid_code = """
erDiagram
    CONCEPT {
        string name
        string description
        string difficulty
    }
    LEARNER {
        string user_id
        string experience_level
        datetime last_studied
    }
    PROGRESS {
        string progress_id
        string user_id
        string concept_name
        string mastery_level
    }
    CONCEPT ||--o{ PROGRESS : tracks
    LEARNER ||--o{ PROGRESS : has
            """
        
        return {
            "topic": topic,
            "mermaid_code": mermaid_code,
            "chart_type": "erDiagram",
            "status": "success"
        }
    
    def _generate_journey_diagram(self, topic: str, topic_lower: str) -> Dict[str, str]:
        """生成旅程图"""
        if "所有权" in topic_lower or "ownership" in topic_lower:
            mermaid_code = """
journey
    title Rust所有权学习之旅
    section 初学阶段
        了解变量所有权: 3: Learner
        编写基础代码: 4: Learner
        遇到所有权错误: 2: Learner
    section 深入理解
        学习转移机制: 5: Learner
        掌握借用检查: 5: Learner
        理解生命周期: 4: Learner
    section 实践应用
        重构代码: 5: Learner
        性能优化: 4: Learner
        内存安全编程: 5: Learner
            """
        elif "trait" in topic_lower or "特征" in topic_lower:
            mermaid_code = """
journey
    title Trait学习与实践
    section 概念理解
        学习Trait概念: 4: Developer
        理解接口定义: 5: Developer
        掌握实现语法: 3: Developer
    section 实践应用
        定义自定义Trait: 5: Developer
        为类型实现Trait: 5: Developer
        使用泛型约束: 4: Developer
    section 高级特性
        理解默认实现: 5: Developer
        使用Trait对象: 4: Developer
        掌握孤儿规则: 3: Developer
            """
        else:
            # 默认旅程图
            mermaid_code = """
journey
    title Rust学习历程
    section 入门阶段
        安装Rust: 5: Learner
        编写Hello World: 5: Learner
        学习基础语法: 4: Learner
    section 进阶阶段
        理解所有权: 3: Learner
        掌握借用和引用: 4: Learner
        学习特征: 5: Learner
    section 高级阶段
        实践项目: 5: Learner
        性能优化: 4: Learner
        深入理解: 5: Learner
            """
        
        return {
            "topic": topic,
            "mermaid_code": mermaid_code,
            "chart_type": "journey",
            "status": "success"
        }
    
    def _detect_chart_type(self, mermaid_code: str) -> str:
        """
        检测Mermaid图表类型
        
        Args:
            mermaid_code: Mermaid代码
            
        Returns:
            图表类型
        """
        code_lower = mermaid_code.lower()
        
        if "flowchart" in code_lower:
            return "flowchart"
        elif "classdiagram" in code_lower:
            return "classDiagram"
        elif "statediagram" in code_lower:
            return "stateDiagram"
        elif "erdiagram" in code_lower:
            return "erDiagram"
        elif "journey" in code_lower:
            return "journey"
        elif "gantt" in code_lower:
            return "gantt"
        elif "pie" in code_lower:
            return "pie"
        elif "gitgraph" in code_lower:
            return "gitgraph"
        else:
            return "flowchart"
    
    def get_visualization_topics(self) -> List[str]:
        """
        获取可可视化的Rust主题列表
        
        Returns:
            Rust主题列表
        """
        return [
            "所有权(Ownership)",
            "借用(Borrowing)",
            "生命周期(Lifetimes)",
            "特征(Traits)",
            "模式匹配(Pattern Matching)",
            "错误处理(Error Handling)",
            "闭包(Closures)",
            "并发编程(Concurrency)",
            "宏(Macros)",
            "迭代器(Iterators)"
        ]
    
    def get_chart_types(self) -> List[str]:
        """
        获取支持的图表类型列表
        
        Returns:
            图表类型列表
        """
        return [
            "流程图",
            "类图",
            "状态图",
            "关系图"
        ]